{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importance Sampling\n",
    "The $J/\\Psi \\rightarrow \\gamma \\pi^0 \\pi^0$ decay is used here (which has a narrow $\\omega$ state)\n",
    "\n",
    "Let's go!\n",
    "\n",
    "### Step 1: Define intensity\n",
    "First we create the intensity. If this step looks unfamiliar to you, you might want to check other intensity construction examples using the expert system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.ui.system_control import (\n",
    "    StateTransitionManager, InteractionTypes)\n",
    "from pycompwa.expertsystem.amplitude.helicitydecay import (\n",
    "    HelicityAmplitudeGeneratorXML)\n",
    "from pycompwa.expertsystem.topology.graph import (\n",
    "    get_intermediate_state_edges)\n",
    "from pycompwa.expertsystem.state.particle import (\n",
    "    get_particle_with_name, particle_list)\n",
    "    \n",
    "initial_state = [(\"J/psi\", [-1, 1])]\n",
    "final_state = [(\"gamma\", [-1, 1]), (\"pi0\", [0]), (\"pi0\", [0])]\n",
    "\n",
    "tbd_manager = StateTransitionManager(initial_state, final_state,\n",
    "                                     formalism_type='helicity',\n",
    "                                     topology_building='isobar')\n",
    "\n",
    "# The omega is so narrow that the hit&miss generation takes way too long.\n",
    "# Therefore we increase the width artificially in this example!\n",
    "omega = get_particle_with_name('omega')\n",
    "parameters = omega['DecayInfo']['Parameter']\n",
    "for par in parameters:\n",
    "    if par['@Type'] == 'Width':\n",
    "        par['Value'] = 0.001\n",
    "\n",
    "tbd_manager.set_allowed_interaction_types(\n",
    "    [InteractionTypes.Strong, InteractionTypes.EM])\n",
    "tbd_manager.allowed_intermediate_particles = ['f2(1270)', 'omega']\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "(solutions, violated_rules) = tbd_manager.find_solutions(\n",
    "        graph_interaction_settings_groups)\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "\n",
    "xml_generator = HelicityAmplitudeGeneratorXML()\n",
    "xml_generator.generate(solutions)\n",
    "xml_generator.write_to_file('model.xml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Create intensity and generate importance weighted phase space sample\n",
    "Now we generate a phase space sample which is importance sampled by the intensity. This can easily be achieved by using the `generate_importance_sampled_phsp()` generate helper function."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    ".. note::\n",
    "   This takes around 30 sec on a Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz running multi-threaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pycompwa is the python interface to ComPWA's c++ modules\n",
    "import pycompwa.ui as pwa\n",
    "\n",
    "# create intensity and kinematics\n",
    "intensity, kin = pwa.create_intensity_and_kinematics('model.xml')\n",
    "\n",
    "# generate phase space sample (flat) used for amplitude normalization\n",
    "gen = pwa.EvtGenGenerator(kin.get_particle_state_transition_kinematics_info())\n",
    "\n",
    "rand_gen = pwa.StdUniformRealGenerator(12345)\n",
    "\n",
    "# generate importance sampled phase space sample\n",
    "phsp_sample_importance = pwa.generate_importance_sampled_phsp(2000, kin, gen, intensity, rand_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Visualize the phase space sample\n",
    "This phase space sample can be visualized in a Dalitz plot. There one expects more events, where the intensity is large, but still an overall uniform distribution. First create all subsystems. Since in this example, we have both $f$ and $\\omega$ resonances, all subsystems already exist and this step is redundant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kin.create_all_subsystems()\n",
    "dataset = pwa.convert_events_to_dataset(phsp_sample_importance, kin)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we use the `create_data_array()` function to get data, which is ready to be put into a numpy record array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.plotting import (\n",
    "    make_dalitz_plots, PlotData, create_nprecord\n",
    ")\n",
    "\n",
    "# use the direct data point access\n",
    "var_names, dataarray = pwa.create_data_array(dataset)\n",
    "data_record = create_nprecord(var_names, dataarray)\n",
    "\n",
    "print(var_names)\n",
    "# check that the sum of weights is equal to the number of events (should be 1000)!\n",
    "print(\"sum of weights (should be 2000):\", sum(data_record.weight))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we create a `PlotData` object and pass it to the `make_dalitz_plots()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_data = PlotData(data_record=data_record)\n",
    "#plotdata.particle_id_to_name_mapping = data_points.get_finalstate_id_to_name_mapping()\n",
    "# plot a 2d histogram\n",
    "make_dalitz_plots(plot_data, [\"mSq_3_4_vs_2\", \"mSq_2_4_vs_3\"], bins=50)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
